from typing import Union

from django.db import transaction
from langchain_chroma import Chroma
from openai.types.beta import VectorStore

import users
from llm.models import VectorStores
from users.models import Users


class VectorStoreDAO:
    @classmethod
    @transaction.atomic
    def add(cls, vector_store: Union[VectorStore,Chroma], user: Users) -> VectorStores:
        """
        添加向量数据库并和用户绑定，同时使其设计适应openai Assistant框架中Vector store
        其中的对应关系如下：

        vector store id->chroma collection name(采取'vs_' +UUID，并开启重复检查)

        vector store name->chroma metadata name字段
        :param vector_store:
        :param user:
        :return:
        """
        if isinstance(vector_store, VectorStore):
            return VectorStores.objects.create(
                vector_store_name=vector_store.name,
                information=vector_store.metadata,
                vector_store_id=vector_store.id,
                user_id=user
            )
        elif isinstance(vector_store, Chroma):
            vs_name=vector_store._collection_metadata["name"]
            return VectorStores.objects.create(
                vector_store_name=vs_name,
                information=vector_store._collection_metadata,
                vector_store_id=vector_store._collection_name,
                user_id=user
            )


